﻿using System;
using System.Diagnostics;
using System.Drawing;
using System.Drawing.Imaging;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;
using System.Threading;
using System.Windows;
using System.Windows.Controls;
using System.Windows.Media;
using System.Windows.Threading;
using DirectShowLib;

namespace MediaBridge
{
    /// <summary>
    /// Method type for notification that the VideoRendererElement has initialized.
    /// </summary>
    public delegate void RendererInitializedDelegate();

    /// <summary>
    /// The VideoRendererElement allows allows writing media samples in WPF at MediaElement speed.
    /// </summary>
    public class VideoRendererElement: FrameworkElement, IMediaBridgeCallback
    {
        #region Private Class Levels
        /// <summary>
        /// P/Invoke for copying native memory.
        /// </summary>
        /// <param name="Destination">Destination of the copy.</param>
        /// <param name="Source">Source of the memory to copy.</param>
        /// <param name="Length">How many bytes of memory to copy.</param>
        [DllImport("kernel32.dll", EntryPoint = "RtlMoveMemory")]
        private static extern void CopyMemory(IntPtr Destination, IntPtr Source, long Length);

        /// <summary>
        /// Internal MediaElement used for rendering.
        /// </summary>
        private MediaElement m_mediaElement;

        /// <summary>
        /// Protocol handler to the DirectShow filter.
        /// </summary>
        private const string DSHOW_PROTOCOL = "DShowMediaBridge://";

        /// <summary>
        /// Used as a unique id so the media bridge 
        /// knows what callback to make.
        /// </summary>
        private readonly Guid m_uniqueIdentifier = Guid.NewGuid();

        /// <summary>
        /// Fires when the element is ready to accept new samples.
        /// </summary>
        public event RendererInitializedDelegate RendererInitialized;

        /// <summary>
        /// Video pixel width of the samples.
        /// </summary>
        private int m_videoPixelWidth;

        /// <summary>
        /// Video pixel height of the samples.
        /// </summary>
        private int m_videoPixelHeight;

        /// <summary>
        /// Reference to the DirectShow filter, COM object.
        /// </summary>
        private IMediaBridgeSourceFilter m_sourceFilter;

        /// <summary>
        /// Used to count the number of samples processed and set
        /// the sample times sent to the renderer.
        /// </summary>
        private int m_frameCounter = 0;

        /// <summary>
        /// The size of the new sample to be written to the renderer.
        /// </summary>
        private int m_newSampleSize = 0;

        /// <summary>
        /// The pointer to the new sample data.
        /// </summary>
        private IntPtr m_currentSamplePointer = IntPtr.Zero;

        /// <summary>
        /// The thread block for the write of the buffer.
        /// </summary>
        private readonly ManualResetEvent m_writeReset = new ManualResetEvent(false);

        /// <summary>
        /// The thread block for the filter callback.
        /// </summary>
        private readonly ManualResetEvent m_fillBufferReset = new ManualResetEvent(false);

#if DEBUG
        /// <summary>
        /// For FPS calculation
        /// </summary>
        private static int m_timesRendered = 0;

        /// <summary>
        /// For FPS calculation
        /// </summary>
        private static DateTime m_recentTime;

        /// <summary>
        /// Used to add internal MediaElement's graph to the
        /// ROT to be viewed in graphedit
        /// </summary>
        private DsROTEntry m_rotEntry;
#endif

        /// <summary>
        /// Flag to determine if we have been initialized.
        /// </summary>
        private bool m_isRendererInitialized = false;

        /// <summary>
        /// Flag to determine if we should tear down the renderer.
        /// </summary>
        private bool m_tearDown = false;

        /// <summary>
        /// The FOURCC code of the MediaSamples.
        /// </summary>
        private int m_fourCC = 0;

        /// <summary>
        /// The MediaSubType for DirectShow.
        /// </summary>
        private Guid m_rendererMediaType;

        /// <summary>
        /// The bits-per-pixel of the samples being rendered.
        /// </summary>
        private int m_bitsPerPixel;

        /// <summary>
        /// Object reference for thread saftey.
        /// </summary>
        private readonly object m_lock = new object();

        /// <summary>
        /// Flags the sample to be a keyframe.
        /// </summary>
        private bool m_isKeyFrame = true;

        #endregion

        #region Overrides
        /// <summary>
        /// When overridden in a derived class, measures the size in layout required for child elements and determines a size for the System.Windows.FrameworkElement-derived class.
        /// </summary>
        /// <param name="availableSize">The available size that this element can give to child elements. Infinity can be specified as a value to indicate that the element will size to whatever content is available.</param>
        /// <returns>The size that this element determines it needs during layout, based on its calculations of child element sizes.</returns>
        protected override System.Windows.Size MeasureOverride(System.Windows.Size availableSize)
        {
            m_mediaElement.Measure(availableSize);
            return base.MeasureOverride(availableSize);
        }

        /// <summary>
        /// When overridden in a derived class, positions child elements and determines a size for a System.Windows.FrameworkElement derived class.
        /// </summary>
        /// <param name="finalSize">The final area within the parent that this element should use to arrange itself and its children.</param>
        /// <returns>The actual size used.</returns>
        protected override System.Windows.Size ArrangeOverride(System.Windows.Size finalSize)
        {
            m_mediaElement.Arrange(new Rect(finalSize));
            return base.ArrangeOverride(finalSize);
        }

        /// <summary>
        /// Overrides System.Windows.Media.Visual.GetVisualChild(System.Int32), and returns a child at the specified index from a collection of child elements.
        /// </summary>
        /// <param name="Index">The zero-based index of the requested child element in the collection.</param>
        /// <returns>The requested child element. This should not return null; if the provided index is out of range, an exception is thrown.</returns>
        protected override Visual GetVisualChild(int Index)
        {
            if (m_mediaElement != null && Index == 0)
            {
                return m_mediaElement;
            }
            else
            {
                throw new ArgumentOutOfRangeException("Index");
            }
        }

        /// <summary>
        /// Gets the number of visual child elements within this element.
        /// </summary>
        protected override int VisualChildrenCount
        {
            get
            {
                return m_mediaElement != null ? 1 : 0;
            }
        }
        #endregion

        #region Dependency Properties
        /// <summary>
        /// Gets or sets a System.Windows.Media.Stretch value that describes how a VideoRendererElement fills the destination rectangle. This is a dependency property.
        /// </summary>
        public Stretch Stretch
        {
            get { return (Stretch)GetValue(StretchProperty); }
            set { SetValue(StretchProperty, value); }
        }

        public static readonly DependencyProperty StretchProperty =
            DependencyProperty.Register("Stretch", 
                                        typeof(Stretch), 
                                        typeof(VideoRendererElement), 
                                        new UIPropertyMetadata(Stretch.None, StretchPropertyChanged));

        private static void StretchPropertyChanged(DependencyObject Sender, DependencyPropertyChangedEventArgs e)
        {
            VideoRendererElement element = Sender as VideoRendererElement;

            if (element == null)
                return;
            
            element.m_mediaElement.Stretch = (Stretch) e.NewValue;
        }

        /// <summary>
        /// Gets or sets a value that determines the restrictions on scaling that are applied to the image. This is a dependency property.
        /// </summary>
        public StretchDirection StretchDirection
        {
            get { return (StretchDirection)GetValue(StretchDirectionProperty); }
            set { SetValue(StretchDirectionProperty, value); }
        }

        public static readonly DependencyProperty StretchDirectionProperty =
            DependencyProperty.Register("StretchDirection", 
                                        typeof(StretchDirection), 
                                        typeof(VideoRendererElement), 
                                        new UIPropertyMetadata(StretchDirection.Both, StretchDirectionPropertyChanged));

        private static void StretchDirectionPropertyChanged(DependencyObject Sender, DependencyPropertyChangedEventArgs e)
        {
            VideoRendererElement element = Sender as VideoRendererElement;

            if (element == null)
                return;

            element.m_mediaElement.StretchDirection = (StretchDirection)e.NewValue;
        }
        #endregion

        /// <summary>
        /// Makes a FOURCC DWORD.
        /// </summary>
        /// <param name="ch0">First character of the FOURCC.</param>
        /// <param name="ch1">Second character of the FOURCC.</param>
        /// <param name="ch2">Third character of the FOURCC.</param>
        /// <param name="ch3">Fourth character of the FOURCC.</param>
        /// <returns>DWORD of the FOURCC generated.</returns>
        public static int MakeFourCC(int ch0, int ch1, int ch2, int ch3)
        {
            return ((byte)(ch0) | ((byte)(ch1) << 8) | ((byte)(ch2) << 16) | ((byte)(ch3) << 24));
        }

        public VideoRendererElement()
        {
            CreateMediaElement();
        }

        /// <summary>
        /// Creates the internal MediaElement we use.
        /// </summary>
        private void CreateMediaElement()
        {
            if (m_mediaElement != null)
                return;

            m_mediaElement = new MediaElement();

            /* Needed so the MediaElement behaviors won't
             * interfere with the VideoRendererElement */
            m_mediaElement.LoadedBehavior = MediaState.Manual;
            m_mediaElement.UnloadedBehavior = MediaState.Manual;

            /* Set the default settings of the video stretch */
            m_mediaElement.Stretch = (Stretch)StretchProperty.DefaultMetadata.DefaultValue;
            m_mediaElement.StretchDirection = (StretchDirection)StretchDirectionProperty.DefaultMetadata.DefaultValue;

            /* Add the MediaElement to the visual tree */
            AddVisualChild(m_mediaElement);
        }

        /// <summary>
        /// Flag to determine if we have been initialized.
        /// </summary>
        public bool IsRendererInitialized
        {
            get { return m_isRendererInitialized; }
        }

        /// <summary>
        /// Writes a new sample to the renderer.
        /// </summary>
        /// <param name="Source">The source pointer to the sample data.</param>
        /// <param name="BufferSize">The size of the buffer.</param>
        /// <param name="IsKeyFrame">The sample is a keyframe.  Uncompressed and interframe (I-Frame) samples should be set to true.</param>
        [MethodImpl(MethodImplOptions.Synchronized)]
        public void WriteSample(IntPtr Source, int BufferSize, bool IsKeyFrame)
        {
            bool shouldReturn = false;

            lock(m_lock)
            {
                if (Source == IntPtr.Zero || 
                    m_currentSamplePointer == IntPtr.Zero || 
                    IsRendererInitialized == false)
                {
                    shouldReturn = true;
                }                
            }

            if (shouldReturn)
                return;

            m_isKeyFrame = IsKeyFrame;

            lock (m_lock)
            {
                /* Copy the source buffer to the destination */
                CopyMemory(m_currentSamplePointer, Source, BufferSize);

                /* Store the size of the sample */
                m_newSampleSize = BufferSize;
            }

            /* Wake up the thread to process the new sample */
            m_fillBufferReset.Set();

            /* Block the current thread until the new sample has been processed */
            m_writeReset.WaitOne();
            m_writeReset.Reset();
        }

        /// <summary>
        /// Writes a Bitmap image to the renderer.  The whole area of the bitmap will be rendered.
        /// </summary>
        /// <param name="BitmapImage">Bitmap image to render.</param>
        public void WriteSample(Bitmap BitmapImage)
        {
            /* Rect full size of bitmap */
            Rectangle imageRect = new Rectangle(0, 0, BitmapImage.Width, BitmapImage.Height);

            WriteSample(BitmapImage, imageRect);
        }

        /// <summary>
        /// Writes a Bitmap image to the renderer.  Only the area specified will be rendered.
        /// </summary>
        /// <param name="BitmapImage">Bitmap image to render.</param>
        /// <param name="PixelRect">Area of the bitmap to render.</param>
        public void WriteSample(Bitmap BitmapImage, Rectangle PixelRect)
        {
            /* Rect full size of bitmap */
            Rectangle imageRect = PixelRect;

            /* Lock the bitmap so we can look up it's skirt */
            BitmapData bmpData = BitmapImage.LockBits(imageRect,
                                                      ImageLockMode.ReadWrite,
                                                      BitmapImage.PixelFormat);

            /* Get the bitmap pixel pointer */
            IntPtr pixelPtr = bmpData.Scan0;

            /* Calculate the buffer size */
            int bufferSize = bmpData.Stride * imageRect.Height;

            /* Write to renderer - uncompressed is always a keyframe*/
            WriteSample(pixelPtr, bufferSize, true);

            /* Unlock pixel bytes */
            BitmapImage.UnlockBits(bmpData);
        }

        /// <summary>
        /// Writes a managed byte array to the renderer.
        /// </summary>
        /// <param name="SampleBuffer">The byte array to render.</param>
        /// <param name="IsKeyFrame">The sample is a keyframe. Uncompressed and interframe (I-Frame) samples should be set to true.</param>
        public void WriteSample(byte[] SampleBuffer, bool IsKeyFrame)
        {
            /* TODO: Instead of de/allocating unmanaged memory
             * for every sample, Marshal.Copy() directly to DShow buffer */

            /* Point to sample unmanaged bytes */
            IntPtr pSampleBytes;

            int byteCount = SampleBuffer.Length;

            /* Alloc some unmanaged bytes */
            pSampleBytes = Marshal.AllocHGlobal(byteCount);

            /* Copy managed bytes to unmanaged pointer */
            Marshal.Copy(SampleBuffer, 0, pSampleBytes, byteCount);

            m_isKeyFrame = IsKeyFrame;

            /* Write the buffer to the renderer */
            WriteSample(pSampleBytes, byteCount, IsKeyFrame);

            /* Free unmanaged bytes */
            Marshal.FreeHGlobal(pSampleBytes);
        }

        /// <summary>
        /// Initializes the renderer.
        /// </summary>
        /// <param name="VideoPixelWidth">The pixel width of the video that will be processed by the renderer.</param>
        /// <param name="VideoPixelHeight">The pixel height of the video that will be processed by the renderer.</param>
        /// <param name="BitsPerPixel">How many bits per pixel the video is.</param>
        /// <param name="FourCC">The four character codec of the video, ie XVID, MJPEG, DIVX.</param>
        /// <param name="MediaType">The GUID of the MediaSubType.  If this is a blank GUID, the DirectShow filter will make the media type from the FourCC value.</param>
        public void InitializeVideo(int VideoPixelWidth, int VideoPixelHeight, int BitsPerPixel, int FourCC, Guid MediaType)
        {
            /* Store the pixel geometry */
            m_videoPixelHeight = VideoPixelHeight;
            m_videoPixelWidth = VideoPixelWidth;

            m_fourCC = FourCC;
            m_bitsPerPixel = BitsPerPixel;
            m_rendererMediaType = MediaType;

            InitializeDirectShow();
        }

        /// <summary>
        /// Initializes the renderer
        /// </summary>
        /// <param name="VideoPixelWidth">The pixel width of the video that will be processed by the renderer.</param>
        /// <param name="VideoPixelHeight">The pixel height of the video that will be processed by the renderer.</param>
        /// <param name="BitsPerPixel">How many bits per pixel the video is.</param>
        /// <param name="FourCC">The four character codec of the video, ie XVID, MJPEG, DIVX.  The static helper function MakeFourCC can generate this value from a string.</param>
        public void InitializeVideo(int VideoPixelWidth, int VideoPixelHeight, int BitsPerPixel, int FourCC)
        {
            /* Store the pixel geometry */
            m_videoPixelHeight = VideoPixelHeight;
            m_videoPixelWidth = VideoPixelWidth;

            m_fourCC = FourCC;
            m_bitsPerPixel = BitsPerPixel;

            /* If the GUID is blank (00000000-0000-0000-0000-000000000000) 
             * the filter will build the media sub-type off the FOURCC */
            m_rendererMediaType = new Guid();

            InitializeDirectShow();
        }

        /// <summary>
        /// Initializes the renderer.
        /// </summary>
        /// <param name="VideoPixelWidth">The pixel width of the video that will be processed by the renderer.</param>
        /// <param name="VideoPixelHeight">The pixel height of the video that will be processed by the renderer.</param>
        /// <param name="BitsPerPixel">How many bits per pixel the video is.</param>
        /// <param name="FourCC">The four character codec of the video, ie XVID, MJPEG, DIVX.</param>
        public void InitializeVideo(int VideoPixelWidth, int VideoPixelHeight, int BitsPerPixel, string FourCC)
        {
            if(FourCC.Length > 4)
            {
                throw new Exception("The FourCC must be 4 characters or less.");
            }

            /* Make FOURCC  */
            int[] fourcc = {0, 0, 0, 0};

            for(int i = 0; i < fourcc.Length; i++)
            {
                fourcc[i] = FourCC[i];
            }
            
            m_fourCC = MakeFourCC(fourcc[0], fourcc[1], fourcc[2], fourcc[3]);

            /* Store the pixel geometry */
            m_videoPixelHeight = VideoPixelHeight;
            m_videoPixelWidth = VideoPixelWidth;
                       
            m_bitsPerPixel = BitsPerPixel;

            /* If the GUID is blank the filter will build 
             * the subtype off the FOURCC */
            m_rendererMediaType = new Guid();

            InitializeDirectShow();
        }

        /// <summary>
        /// Configures the renderer to use uncompressed RGB samples.
        /// </summary>
        /// <param name="VideoPixelWidth">The pixel width of the video that will be processed by the renderer.</param>
        /// <param name="VideoPixelHeight">The pixel height of the video that will be processed by the renderer.</param>
        /// <param name="BitsPerPixel">Options are 1, 4, 24 and 32 bits per pixel.</param>
        public void InitializeVideo(int VideoPixelWidth, int VideoPixelHeight, int BitsPerPixel)
        {
            switch(BitsPerPixel)
            {
                case 1:
                    m_rendererMediaType = RendererMediaType.RGB1;
                    break;
                case 4:
                    m_rendererMediaType = RendererMediaType.RGB4;
                    break;
                case 32:
                    m_rendererMediaType = RendererMediaType.RGB32;
                    break;
                case 24:
                    m_rendererMediaType = RendererMediaType.RGB24;
                    break;
                default:
                    throw new Exception("Bits per pixel options are 1, 4, 24 and 32.");
            }

            /* Store the pixel geometry */
            m_videoPixelHeight = VideoPixelHeight;
            m_videoPixelWidth = VideoPixelWidth;

            m_fourCC = (int)BitmapCompression.UncompressedRGB;

            m_bitsPerPixel = BitsPerPixel;

            InitializeDirectShow();
        }

        /// <summary>
        /// Registers a callback and loads up the DirectShow bridge filter.
        /// </summary>
        private void InitializeDirectShow()
        {
            m_fillBufferReset.Set();
            m_writeReset.Set();

            m_tearDown = false;

            /* Build the Url to feed the MediaElement */
            string url = DSHOW_PROTOCOL + m_uniqueIdentifier;

            /* Register the callback for the DirectShow filter */
            MediaBridgeManager.RegisterCallback(url, BridgeCallback);

            /* Set the source on MediaElement */
            m_mediaElement.Source = new Uri(url);

            /* Make the magic happen - The callback
             * should execute next on another thread 
             * created by WMP OCX under the covers */
            m_mediaElement.Play();
        }

        /// <summary>
        /// Callback from the DirectShow filter.  This happens from it's own thread.
        /// </summary>
        /// <param name="GraphInfo">Information passed back from the DirectShow filter.</param>
        private void BridgeCallback(MediaBridgeGraphInfo GraphInfo)
        {
            /* Cast the pointer to a COM object */
            m_sourceFilter = (IMediaBridgeSourceFilter)Marshal.GetObjectForIUnknown(GraphInfo.SourceFilter);

#if DEBUG
            /* Add graph to the ROT so we can view the graph in graghedt.exe */
            IFilterGraph graph = (IFilterGraph)Marshal.GetObjectForIUnknown(GraphInfo.FilterGraph);
            m_rotEntry = new DsROTEntry(graph);
            Marshal.ReleaseComObject(graph);
            graph = null;

            /* For FPS count */
            m_recentTime = DateTime.Now;
#endif

            /* Call a COM method to set the width/height of the video */
            m_sourceFilter.InitializeVideo(m_videoPixelWidth, 
                                           m_videoPixelHeight, 
                                           m_bitsPerPixel, 
                                           m_fourCC, 
                                           m_rendererMediaType);

            /* Tell the filter to make callbacks to our object */
            m_sourceFilter.SetCallback(this);

            /* Release our COM reference - Manual memory management is a bitch */
            if (m_sourceFilter != null)
                Marshal.ReleaseComObject(m_sourceFilter);

            m_sourceFilter = null;

            /* Notify that we are ready to get samples */
            FireRendererInitializedEvent();
        }

        /// <summary>
        /// Fires an event to notify the VideoRenderer
        /// is ready to accept samples.
        /// </summary>
        private void FireRendererInitializedEvent()
        {
            /* Run via the Dispatcher */
            Dispatcher.Invoke(DispatcherPriority.Normal, (ThreadStart)delegate
            {
                m_isRendererInitialized = true;
                if (RendererInitialized != null)
                    RendererInitialized();         
            });
        }

        /// <summary>
        /// Cleans up the VideoRendererElement for 
        /// disposal or reinitialization.
        /// </summary>
        public void TearDown()
        {
            m_tearDown = true;
            m_frameCounter = 0;

#if DEBUG
            m_timesRendered = 0;

            if(m_rotEntry != null)
                m_rotEntry.Dispose();
#endif
            Dispatcher.Invoke(DispatcherPriority.Normal, (ThreadStart)delegate
            {
                /* Wake up the threads */
                m_fillBufferReset.Set();
                m_writeReset.Set();

                /* Shut down the MediaElement - Free ish up */
                m_mediaElement.Close();
            });
        }

        /// <summary>
        /// Is called from DirectShow filter after each frame has been rendered.
        /// </summary>
        /// <param name="pSample">The media sample COM object.</param>
        /// <returns>HRESULT of the method.</returns>
        int IMediaBridgeCallback.FillBuffer(IMediaSample pSample)
        {
            m_frameCounter++;

            if (m_tearDown)
            {
                m_isRendererInitialized = false;
                pSample.SetActualDataLength(0);

                /* Run COMobj.Release() on COM */
                Marshal.ReleaseComObject(pSample);
                return 2;
            }

            lock(m_lock)
            {
                pSample.GetPointer(out m_currentSamplePointer);
            }
            
            /* Wake up the thread writing to the renderer */
            m_writeReset.Set();
           
            /* Block this thread */
            m_fillBufferReset.WaitOne();
            m_fillBufferReset.Reset();
            
            lock(m_lock)
            {
                /* Set the buffer size on the sample COM object */
                pSample.SetActualDataLength(m_newSampleSize);

                /* Reset variables */
                m_newSampleSize = 0;
                m_currentSamplePointer = IntPtr.Zero;
            }
            
            /* Configure a sample time to execute immediately 
             * without blocks from the DShow renderer */
            pSample.SetTime(m_frameCounter, m_frameCounter + 1);
                        
            pSample.SetSyncPoint(m_isKeyFrame);

#if DEBUG
            m_timesRendered++;

            /* Do FPS calculation */
            if ((DateTime.Now - m_recentTime) > TimeSpan.FromSeconds(1))
            {
                Debug.WriteLine("VideoRendererFPS: " + m_timesRendered);
                m_timesRendered = 0;
                m_recentTime = DateTime.Now;
            }
#endif
            /* Run comObj.Release() on COM */
            Marshal.ReleaseComObject(pSample);

            /* Return S_OK */
            return 0;
        }
    }
}
